import os
import sys
from datetime import datetime
from PyQt5.QtGui import QIcon, QImage
from PyQt5.QtWidgets import QSystemTrayIcon, QMenuBar, QAction, QHBoxLayout
from PyQt5.QtWidgets import QApplication, QMainWindow, QLineEdit, QTextEdit, QVBoxLayout, QWidget, QTabWidget, QDialog, QPushButton
import yaml
from transformers import TextIteratorStreamer
from unsloth import FastLanguageModel
####################################################################################

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
with open(os.path.join(BASE_DIR, 'conf.yaml'), 'r', encoding='utf-8') as f:
    config = yaml.safe_load(f)
####################################################################################
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name= config['model'],  # 加载训练后的LoRA模型
    max_seq_length=2048,
    dtype=None,
    load_in_4bit=True,
)


def llama3_response(request):
    _alpaca_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
        ### Instruction:
        {}
        ### Input:
        {}
        ### Response:
        {}"""
    _text = _alpaca_prompt.format(
        "请用中文回答",
        request,
        "",
    )
    _inputs = tokenizer([_text, ], return_tensors="pt")
    _streamer = TextIteratorStreamer(tokenizer=tokenizer)
    model.generate(**_inputs, streamer=_streamer, max_new_tokens=128, pad_token_id=tokenizer.eos_token_id)
    _response = ""
    for _token in _streamer:
        _response = _response + _token
    # _response = _response+"\n"
    # _index = _response.find("### Response:")+len("### Response:")
    # if _index != -1:
    #     return _response[_index:-1 - len("<|end_of_text|>")]
    # else:
    #     return "\n"
    _response= _response[len(_text):-1 - len("<|end_of_text|>")]
    if "请用中文回答" in _response:
        start_index = _response.index("请用中文回答")
        if start_index != len(_response) - len("请用中文回答"):
            return _response[start_index + len("请用中文回答"):]
        else:
            return ""  # 或者其他适当的默认值
    return ""

FastLanguageModel.for_inference(model)


class Chat(QMainWindow):
    def __init__(self):
        super(Chat, self).__init__()
        self.desktop = QApplication.desktop()
        self.screenRect = self.desktop.screenGeometry()
        self.screenheight, self.screenwidth = self.screenRect.height(), self.screenRect.width()
        self.img = QImage().load(os.path.join(BASE_DIR, config['icon']))
        self.setWindowIcon(QIcon(config['icon']))
        self.tp = QSystemTrayIcon(self)  # 初始化系统托盘
        self.tp.setIcon(QIcon((config['icon'])))
        self.tp.setToolTip('llama3')
        # def init_window(self):
        """
        初始化窗口
        """
        self.is_follow_mouse = False  # 初始化鼠标没有移动
        self.mouse_drag_pos = self.pos()

        self.setWindowTitle(config['model'])
        # self.setGeometry(100, 100, 400, 300)
        # self.setWindowState(window.windowState() | Qt.WindowFullScreen)
        # self.showFullScreen()

        # 设置窗口的最小大小
        self.setMinimumSize(1080,720)

        # 设置窗口的最大大小
        # self.setMaximumSize(2560,1920)
        self.adjustSize()

        # 初始化tab_widget
        self.tab_widget = QTabWidget(self)
        self.setCentralWidget(self.tab_widget)


        # 创建第一个标签页
        self.tab1 = QWidget()
        self.tab_widget.addTab(self.tab1, "AI对话")

        # 创建第二个标签页
        self.tab2 = QWidget()
        self.tab_widget.addTab(self.tab2, "设置")

        # 创建布局
        self.layout_chat = QVBoxLayout(self.tab1)
        # 创建输出框
        self.output_box = QTextEdit(self.tab1)
        # 创建输入框
        self.input_box = QTextEdit(self.tab1)
        # 为每个标签页添加布局和控件
        self.setup_tab_chat(self.tab1)
        self.setup_tab_setup(self.tab2)
        self.show()  # 显示窗口
    # def tray(self):
        """
        建立一个托盘
        """
        self.bar = QMenuBar(self)
        self.menu = self.bar.addMenu('菜单')
        quit = QAction('退出', self)
        quit.triggered.connect(self.quit)
        self.menu.addAction(quit)
        self.tp.setContextMenu(self.menu)
        self.tp.show()

    def setup_tab_chat(self, tab):

        self.output_box.setPlaceholderText("输出内容")
        self.output_box.setReadOnly(True)
        self.layout_chat.addWidget(self.output_box)

        self.input_box.setPlaceholderText("输入内容")
        self.input_box.setFixedHeight(5 * 30)
        # 创建确认按钮
        confirm_button = QPushButton("确认", tab)
        confirm_button.clicked.connect(self.on_confirm)
        # 创建水平布局并将输入框和按钮添加进去
        hbox = QHBoxLayout()
        hbox.addWidget(self.input_box)
        hbox.addWidget(confirm_button)

        # 将水平布局添加到垂直布局
        self.layout_chat.addLayout(hbox)

    def setup_tab_setup(self, tab):
        # 创建布局
        layout = QVBoxLayout(tab)

        # 创建输出框
        output_box = QTextEdit(tab)
        output_box.setPlaceholderText("设置")
        output_box.setReadOnly(True)
        layout.addWidget(output_box)
    def on_confirm(self):
        # 这里可以添加处理输入的逻辑
        text = datetime.now().strftime("%Y年%m月%d日 %H时%M分: \n")+self.input_box.toPlainText()+'\n'
        self.output_box.append(text)  # 将输入框的内容追加到输出框
        text = datetime.now().strftime("%Y年%m月%d日 %H时%M分: \n") + llama3_response(self.input_box.toPlainText())+'\n'
        self.output_box.append(text)  # 将输入框的内容追加到输出框
    def quit(self):
        sys.exit()

if __name__ == '__main__':
    # 创建程序和对象
    app = QApplication(sys.argv)
    Chat = Chat()
    sys.exit(app.exec_())
